"""
FILE: scan_mesher.py

LAST MODIFIED: 25-Sept-2009

DESCRIPTION:
Class for creating a fieldwork mesh over a stack of images
"""

import scipy
import sys
sys.path.append( '../' )
from field import template_fields
from field.topology import element_types

class FemurMesher:
	""" class for meshing femurs
	"""
	
	def __init__( self, scan, geometric_field ):
		self.scan = scan
		self.GF = geometric_field
		
		self.shaft = None	# shaft subvolume
	
	#==================================================================#
	def find_shaft( self, stdTol = 0.5 ):
		""" locates the shaft of the femur and instantiate a subvolume
		
		starting from the middle slice in the stack z direction 
		(shape[0]), examines every (spacing) slice and calculates the
		radius stdev. If a slice's radius stdev exceeds tolerance from
		the cumulative mean of previous slices, it is deemed either the
		top or bottom of the shaft
		"""
		
		spacing = 2
		slice0 = self.scan.I.shape[0]/2
		sliceUp = slice0 - spacing
		sliceDown = slice0 + spacing
		done = 0
		foundTop = 0
		foundBottom = 0
		meanStd = None
		radiusStd = []
		radiusMean = []
		
		# do initial slice
		slice = self.scan.createSubSlice( 0, slice0 )
		sliceR = slice.calculateSlicePolar()[0]
		radiusStd.append( sliceR.std() )
		radiusMean.append( sliceR.mean() )
		meanStd = scipy.mean( radiusStd )
		
		while not done:
			
			if ( not foundTop ) and ( sliceUp >= 0 ):
				slice = self.scan.createSubSlice( 0, sliceUp )
				sliceR = slice.calculateSlicePolar()[0]
				if abs( sliceR.std() - meanStd ) > stdTol:
					foundTop = 1
				else:
					radiusStd.insert( 0, sliceR.std() )
					radiusMean.insert( 0, sliceR.mean() )
					meanStd = scipy.mean( radiusStd )
					sliceUp -= spacing

			if ( not foundBottom ) and ( sliceDown <= (self.scan.I.shape[0] - 1) ):
				slice = self.scan.createSubSlice( 0, sliceDown )
				sliceR = slice.calculateSlicePolar()[0]
				if abs( sliceR.std() - meanStd ) > stdTol:
					foundBottom = 1
				else:
					radiusStd.append( sliceR.std() )
					radiusMean.append( sliceR.mean() )
					meanStd = scipy.mean( radiusStd )
					sliceDown += spacing
		
			if foundTop and foundBottom:
				done = 1
				
		# instantiate shaft subvolume
		self.shaft = self.scan.createSubVolumes( 'shaft', [sliceUp, 0, 0], [sliceDown - sliceUp, self.scan.I.shape[1], self.scan.I.shape[2]] )
		self.shaft_offset = [sliceUp, 0, 0]
		
		print 'shaft slices:', [sliceUp, sliceDown]
		
		return ( self.shaft, [sliceUp, sliceDown], radiusStd, radiusMean, meanStd )
			
	#==================================================================#	
	def mesh_shaft( self, ring_n, ring_element_n):
		""" generate stack of ring meshes for the shaft of a femur
		"""

		points_per_slice = ring_element_n * 2
		points_per_ring = ring_element_n * 6
		
		# get positions of geometric points
		points = scipy.array( self._shaft_get_points( ring_n, ring_element_n ) )
		#~ points = scipy.add( points, offset )
		
		# create rings
		for i in range( ring_n ):
			# get the points for current ring
			if i == 0:
				ring_points = points[ 0 : points_per_ring ]
			else:
				ring_points = points[ i*points_per_ring - i*points_per_slice : (i+1)*points_per_ring - i*points_per_slice ]
			#~ elif i == (ring_n - 1):
				#~ ring_points = points[ i*points_per_ring - points_per_slice : ( i + 1 )*points_per_ring - points_per_slice ]
			#~ else:
				#~ ring_points = points[ i*points_per_ring - (i-1)*points_per_slice : (i+1)*points_per_ring - (i-1)*points_per_slice ]
	#~ 
			if not self._shaft_add_ring( ring_points, ring_element_n ):
				print 'ERROR: FemurMesher.mesh_shaft: failed to add ring', i
				print 'number of ring points:', len( ring_points )
				print 'ring points:', ring_points
				return
				
		return 1
	
	#==================================================================#
	def _shaft_add_ring( self, points, ring_element_n ):
		""" points = array([[x0,y0,z0],[x1,y1,z1],[...],...])
		"""
		
		ring = template_fields.two_quad_ring( ring_element_n )
		
		parameters = []
		for i in range( 3 ):
			parameters.append( [ [v] for v in points[:,i] ] )
			
		if self.GF.add_element_with_parameters( ring, parameters ):
			return 1
		else:
			print 'ERROR: FemurMesher._shaft_add_ring: failed to add ring element'
			return
			
	#==================================================================#
	def _shaft_get_points( self, ring_n, ring_element_n ):
		""" find geometric points along the shaft of the femur
		"""
		
		slice_n = ring_n * 2 + 1	# number of slices
		slice_points_n = ring_element_n * 2		# points per slice
		slices = scipy.linspace( 0, self.shaft.I.shape[0]-1, slice_n ).round().astype( int )	# slice numbers

		points = []
		ridge_n = 0
		#~ slice_point0s = []	# first point in each slice
		for slice in slices:
			s = self.shaft.createSubSlice( 0, slice )
			# get points from this slice( 2D image )
			slice_points = self._shaft_get_slice_points( s, slice_points_n, slice )
			slice_points = scipy.add( slice_points, self.shaft_offset )
			#~ slice_point0s.append( slice_points[0] )
			for p in slice_points:
				points.append( p )
			
			# add 1st point in slice to geometric field as a landmark
			self.GF.add_geometric_point( [ [v] for v in slice_points[0] ], name = 'shaft_ridge_'+str(ridge_n) )
			ridge_n += 1
			
		return points

	#==================================================================#
	def _shaft_get_slice_points( self, slice, nPoints, sliceN):
		""" get_slice_points( slice, nPoints, sliceN ) -> list of 
		(x,y) for each point
		
		slice is a 2D image array
		centre: [x,y] position of the centre of the object
		startV: [v1,v2] vector of the starting position from centre
		nPoints: integer number of points to get 
		"""
		
		points = []
		averages = 10		# take position of closest 10 points
		# starting angle
		startV = slice.pAxes[ :, slice.pAxesMag.argmax() ]
		
		theta0 = calcTheta( startV[0], startV[1] )
		dtheta = 2*scipy.pi / nPoints
		
		print startV
		print theta0
		
		# get all nonzero point indices
		xi, yi = scipy.nonzero( slice.I )
		# center
		xi = xi - slice.CoM[0]
		yi = yi - slice.CoM[1]
		
		# get theta coords of all pixels in slice
		theta = slice.calculateSlicePolar()[1]
		
		for i in range( nPoints ):
			# find closes data theta to t
			t = theta0 + i*dtheta
			if t > 2*scipy.pi:
				t -= 2*scipy.pi
			
			# get closest point indices
			thetaDist = (theta - t)**2
			tempPoints = []
			for j in range( averages ):
				closest_i = thetaDist.argmin()
				thetaDist[ closest_i ] = scipy.inf
				
				closest_c = [ xi[closest_i], yi[closest_i] ]
				closest_c = [sliceN] + list( scipy.add( closest_c, slice.CoM ) )
				tempPoints.append( closest_c )
				
			# average
			tempPoints = scipy.array( tempPoints )
			points.append( tempPoints.mean( axis = 0 ) )
			
		return points
		
#======================================================================#

def calcTheta( x, y ):
	""" angle anticlockwise from the +ve x axis of a point (x,y)
	"""
	
	s = ( scipy.sign(x), scipy.sign(y) )
	arct = scipy.arctan( y/x )
	case = { (1, 1): arct,\
	         (-1, 1): scipy.pi + arct,\
			 (-1, -1): scipy.pi + arct,\
			 (1, -1): 2*scipy.pi + arct }
			 
	return case[s]
